# 指定文件编码为UTF-8
# coding: utf-8

"""
util.py 功能说明：
1. 提供神经网络训练和卷积运算的实用工具函数
2. 主要包含四类功能：
   - 数据平滑处理(smooth_curve)
   - 数据集随机打乱(shuffle_dataset)
   - 卷积运算辅助函数(im2col/col2im)
   - 卷积输出尺寸计算(conv_output_size)

关键功能：
1. im2col/col2im: 实现图像到矩阵的高效转换，加速卷积运算
2. shuffle_dataset: 训练前打乱数据顺序，避免模型学习到输入顺序
3. smooth_curve: 平滑训练曲线，便于可视化分析
4. 支持带padding和stride的卷积运算

典型应用场景：
- 卷积神经网络的前向传播和反向传播
- 训练过程监控和可视化
- 数据预处理
"""

import numpy as np

def smooth_curve(x):
    """使用滑动窗口平滑数据曲线(常用于损失函数可视化)

    参数:
        x: 输入数据序列
    返回:
        平滑后的数据序列

    实现方法:
        1. 使用Kaiser窗口函数进行卷积
        2. 处理边界效应
    """
    window_len = 11  # 滑动窗口长度
    # 扩展数据边界处理边界效应
    s = np.r_[x[window_len-1:0:-1], x, x[-1:-window_len:-1]]
    # 生成Kaiser窗口(β=2)
    w = np.kaiser(window_len, 2)
    # 执行归一化卷积
    y = np.convolve(w/w.sum(), s, mode='valid')
    # 去除边界扩展部分
    return y[5:len(y)-5]

def shuffle_dataset(x, t):
    """随机打乱数据集顺序(保持数据和标签对应关系)

    参数:
        x: 训练数据(可以是2D或4D数组)
        t: 对应标签
    返回:
        打乱顺序后的x和t

    注意:
        支持2D(全连接网络)和4D(卷积网络)输入数据
    """
    # 生成随机排列索引
    permutation = np.random.permutation(x.shape[0])
    # 根据维度处理训练数据
    x = x[permutation,:] if x.ndim == 2 else x[permutation,:,:,:]
    # 同步打乱标签
    t = t[permutation]
    return x, t

def conv_output_size(input_size, filter_size, stride=1, pad=0):
    """计算卷积/池化后的输出尺寸

    公式:
        (输入大小 + 2*pad - 滤波器大小)/stride + 1

    参数:
        input_size: 输入尺寸(高或宽)
        filter_size: 滤波器尺寸
        stride: 步长(默认1)
        pad: 填充(默认0)
    返回:
        输出尺寸
    """
    return (input_size + 2*pad - filter_size) // stride + 1

def im2col(input_data, filter_h, filter_w, stride=1, pad=0):
    """将4D输入数据转换为2D矩阵(用于高效卷积运算)

    参数:
        input_data: 4D输入数据(数量,通道,高,宽)
        filter_h: 滤波器高度
        filter_w: 滤波器宽度
        stride: 步长(默认1)
        pad: 填充(默认0)
    返回:
        2D矩阵

    工作原理:
        将每个局部感受野展平为一行
    """
    N, C, H, W = input_data.shape
    # 计算输出尺寸
    out_h = (H + 2*pad - filter_h) // stride + 1
    out_w = (W + 2*pad - filter_w) // stride + 1

    # 添加padding
    img = np.pad(input_data, [(0,0), (0,0), (pad, pad), (pad, pad)], 'constant')
    # 初始化结果矩阵
    col = np.zeros((N, C, filter_h, filter_w, out_h, out_w))

    # 滑动窗口提取局部区域
    for y in range(filter_h):
        y_max = y + stride*out_h
        for x in range(filter_w):
            x_max = x + stride*out_w
            col[:, :, y, x, :, :] = img[:, :, y:y_max:stride, x:x_max:stride]

    # 调整维度并展平
    col = col.transpose(0, 4, 5, 1, 2, 3).reshape(N*out_h*out_w, -1)
    return col

def col2im(col, input_shape, filter_h, filter_w, stride=1, pad=0):
    """将im2col转换后的矩阵恢复为4D图像数据(用于反向传播)

    参数:
        col: im2col输出的2D矩阵
        input_shape: 原始输入形状(数量,通道,高,宽)
        filter_h: 滤波器高度
        filter_w: 滤波器宽度
        stride: 步长(默认1)
        pad: 填充(默认0)
    返回:
        4D图像数据

    注意:
        这是im2col的逆操作
    """
    N, C, H, W = input_shape
    # 计算输出尺寸
    out_h = (H + 2*pad - filter_h) // stride + 1
    out_w = (W + 2*pad - filter_w) // stride + 1
    # 调整矩阵形状
    col = col.reshape(N, out_h, out_w, C, filter_h, filter_w).transpose(0, 3, 4, 5, 1, 2)

    # 初始化输出图像(考虑padding)
    img = np.zeros((N, C, H + 2*pad + stride - 1, W + 2*pad + stride - 1))
    # 滑动窗口恢复数据
    for y in range(filter_h):
        y_max = y + stride*out_h
        for x in range(filter_w):
            x_max = x + stride*out_w
            img[:, :, y:y_max:stride, x:x_max:stride] += col[:, :, y, x, :, :]

    # 去除padding部分
    return img[:, :, pad:H + pad, pad:W + pad]

"""
使用示例：
1. 卷积运算加速：
   col = im2col(input_data, 3, 3)  # 3x3滤波器
   col_filter = filter.reshape(-1, 1)
   out = np.dot(col, col_filter)
   out = out.reshape(N, out_h, out_w)

2. 数据预处理：
   x_train, t_train = shuffle_dataset(x_train, t_train)

3. 可视化平滑：
   smoothed_loss = smooth_curve(loss_history)

注意事项：
1. im2col会显著增加内存使用，但大幅提升卷积速度
2. shuffle_dataset应在每个epoch前调用
3. 平滑操作会引入滞后，仅用于可视化
"""
